{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Magpie Toy Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "848f65bd3c1243fea1e91865f48489dd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "import torch\n",
    "import json\n",
    "from transformers import AutoTokenizer\n",
    "from fastchat.model import get_conversation_template\n",
    "\n",
    "model_id = \"meta-llama/Meta-Llama-3-8B-Instruct\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, add_special_tokens=True)\n",
    "\n",
    "pipeline = transformers.pipeline(\n",
    "    \"text-generation\",\n",
    "    model=model_id,\n",
    "    model_kwargs={\"torch_dtype\": torch.bfloat16},\n",
    "    device=\"cuda:0\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Let's prepare the left-side template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "with open(\"configs/model_configs.json\", \"r\") as f:\n",
    "    model_configs = json.load(f)\n",
    "    model_config = model_configs[model_id]\n",
    "\n",
    "# Prompt for extracting instructions from Llama-3-8B-Instruct\n",
    "pre_query_template = model_config[\"pre_query_template\"]\n",
    "print(pre_query_template)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 1: Extracting Instructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:128009 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracted Instruction: I'd like to get a new pair of shoes for the summer. What kind of shoes are good for warm weather?\n"
     ]
    }
   ],
   "source": [
    "terminators = [\n",
    "    tokenizer.eos_token_id,\n",
    "    tokenizer.convert_tokens_to_ids(\"<|eot_id|>\"),\n",
    "]\n",
    "\n",
    "instruction = pipeline(\n",
    "    pre_query_template,\n",
    "    max_new_tokens=2048,\n",
    "    eos_token_id=terminators,\n",
    "    do_sample=True,\n",
    "    temperature=1,\n",
    "    top_p=1,\n",
    ")\n",
    "\n",
    "sanitized_instruction = instruction[0]['generated_text'][len(pre_query_template):].split(\"\\n\")[0]\n",
    "print(f\"Extracted Instruction: {sanitized_instruction}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2: Generating Responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:128009 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conversation(name='llama-3', system_template='<|start_header_id|>system<|end_header_id|>\\n\\n{system_message}<|eot_id|>', system_message='', roles=('user', 'assistant'), messages=[], offset=0, sep_style=<SeparatorStyle.LLAMA3: 8>, sep='', sep2=None, stop_str='<|eot_id|>', stop_token_ids=[128001, 128009], max_image_size_mb=None)\n"
     ]
    }
   ],
   "source": [
    "conv = get_conversation_template(\"llama-3\")\n",
    "print(conv) # Make sure the conversation template is correct with name='llama-3' !!!\n",
    "conv.append_message(conv.roles[0], sanitized_instruction)\n",
    "conv.append_message(conv.roles[1], None)\n",
    "prompt = conv.get_prompt()\n",
    "\n",
    "response = pipeline(\n",
    "    prompt,\n",
    "    max_new_tokens=2048,\n",
    "    eos_token_id=terminators,\n",
    "    do_sample=True,\n",
    "    temperature=1,\n",
    "    top_p=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We have successfully created an instruction-response pair!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User's message: I'd like to get a new pair of shoes for the summer. What kind of shoes are good for warm weather?\n",
      "Assitant's response: Perfect timing! When it comes to shoes for warm weather, you'll want something comfortable, breathable, and stylish. Here are some great options:\n",
      "\n",
      "1. **Sandals**: A classic choice for summer, sandals are perfect for walking, beach trips, or lounging around. Look for ones with good arch support and adjustable straps.\n",
      "2. **Sneakers with breathable materials**: Sneakers made from mesh, canvas, or mesh panels are great for hot weather. Brands like New Balance, ASICS, or Converse offer such options.\n",
      "3. **Espadrilles**: These flat, rope-soled shoes are perfect for casual outings. They're lightweight, breathable, and can be dressed up or down.\n",
      "4. **Slide sandals**: Similar to traditional sandals, but with a more relaxed design. They're great for poolside or beach activities.\n",
      "5. **Loafers**: A stylish option for warm weather, loafers are lightweight, comfortable, and can be dressed up or down.\n",
      "6. **Water shoes**: If you plan to spend time in or around water (e.g., swimming, kayaking, or paddleboarding), waterproof or water-friendly shoes are a must-have.\n",
      "7. **Mules**: These backless shoes are perfect for warm weather. Look for ones with a soft, breathable upper and a cushioned insole.\n",
      "\n",
      "When shopping for shoes for warm weather, consider the following features:\n",
      "\n",
      "* Breathable materials (e.g., mesh, canvas, leather with mesh panels)\n",
      "* Lightweight construction\n",
      "* Adjustable straps or closures for a secure fit\n",
      "* Cushioned insoles for added comfort\n",
      "* Water-friendly or quick-drying materials (if you plan to get wet)\n",
      "* A simple, casual design for easy pairing with lightweight outfits\n",
      "\n",
      "Which type of shoes do you think you might be interested in?\n"
     ]
    }
   ],
   "source": [
    "print(f\"User's message: {sanitized_instruction}\")\n",
    "print(f\"Assitant's response: {response[0]['generated_text'][len(prompt):]}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "datareverse",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
